// Copyright (c) 2015-2016, tandasat. All rights reserved.
// Use of this source code is governed by a MIT-style license that can be
// found in the LICENSE file.

/// @file
/// Implements DdiMon functions.

#include "ddi_mon.h"
#include <ntimage.h>
#define NTSTRSAFE_NO_CB_FUNCTIONS
#include <ntstrsafe.h>
#include "../HyperPlatform/common.h"
#include "../HyperPlatform/log.h"
#include "../HyperPlatform/util.h"
#include "../HyperPlatform/ept.h"
#include "../HyperPlatform/kernel_stl.h"
#include "../HyperPlatform/performance.h"
#include "shadow_hook.h"
#include "main.h"
#include "hook.h"
#include "NativeEnums.h"
#include "../../Shared/Protocol.h"

extern CProcList *m_IgnoreProcList;
extern CFileList *m_IgnoreFileList;
extern CEventList *m_EventList;

extern PFLT_FILTER m_pFilterHandle;
extern PFLT_PORT m_pClientPort;

EXTERN_C
{

extern DYNAMIC_DATA dynData;

extern POBJECT_TYPE *PsProcessType;

NTSTATUS GetProcessIdByHandle(__in HANDLE ProcessHandle, __out PHANDLE ProcessId);

}

//NtQuerySystemInformation

NTSTATUS NTAPI NewNtQuerySystemInformation(
	SYSTEM_INFORMATION_CLASS SystemInformationClass,
	PVOID SystemInformation,
	ULONG SystemInformationLength,
	PULONG ReturnLength);

static ShadowHookTarget m_NtQuerySystemInformationHookTarget =
{
	NewNtQuerySystemInformation,
	nullptr,
};

NTSTATUS NTAPI NewNtQuerySystemInformation(
	SYSTEM_INFORMATION_CLASS SystemInformationClass,
	PVOID SystemInformation,
	ULONG SystemInformationLength,
	PULONG ReturnLength)
{
	const auto original = (NTSTATUS(NTAPI*)(SYSTEM_INFORMATION_CLASS, PVOID, ULONG, PULONG))m_NtQuerySystemInformationHookTarget.original_call;
	
	const auto status = original(SystemInformationClass, SystemInformation, SystemInformationLength, ReturnLength);

	svc_nt_query_systeminfo_data *data = NULL;

	if (m_EventList->IsCapturing())
	{
		if (NT_SUCCESS(status) && ExGetPreviousMode() == UserMode)
		{
			data = (svc_nt_query_systeminfo_data *)
				ExAllocatePoolWithTag(PagedPool, sizeof(svc_nt_query_systeminfo_data), 'TXSB');
			if (data)
			{
				RtlZeroMemory(data, sizeof(svc_nt_query_systeminfo_data));
				data->protocol = svc_nt_query_systeminfo;
				data->size = sizeof(svc_nt_query_systeminfo_data);
				data->time = PerfGetSystemTime();
				data->ProcessId = (ULONG)PsGetCurrentProcessId();
				data->QueryClass = (ULONG)SystemInformationClass;

				m_EventList->SendEvent(data);
			}
		}
	}

	return status;
}

//NtOpenProcess

NTSTATUS NTAPI NewNtOpenProcess(
	_Out_    PHANDLE            ProcessHandle,
	_In_     ACCESS_MASK        DesiredAccess,
	_In_     POBJECT_ATTRIBUTES ObjectAttributes,
	_In_opt_ PCLIENT_ID         ClientId
);

static ShadowHookTarget m_NtOpenProcessHookTarget =
{
	NewNtOpenProcess,
	nullptr,
};

NTSTATUS NTAPI NewNtOpenProcess(
	_Out_    PHANDLE            ProcessHandle,
	_In_     ACCESS_MASK        DesiredAccess,
	_In_     POBJECT_ATTRIBUTES ObjectAttributes,
	_In_opt_ PCLIENT_ID         ClientId
)
{
	const auto original = (NTSTATUS(NTAPI *)(PHANDLE, ACCESS_MASK, POBJECT_ATTRIBUTES, PCLIENT_ID))m_NtOpenProcessHookTarget.original_call;

	NTSTATUS status = original(ProcessHandle, DesiredAccess, ObjectAttributes, ClientId);

	svc_nt_open_process_data *data = NULL;

	if (m_EventList->IsCapturing())
	{
		__try
		{
			CLIENT_ID CapturedCid;
			if (ARGUMENT_PRESENT(ClientId) && ExGetPreviousMode() == UserMode)
			{
				ProbeForRead(ClientId, sizeof(CLIENT_ID), sizeof(ULONG));
				CapturedCid = *ClientId;

				PEPROCESS Process = NULL;
				if (CapturedCid.UniqueProcess != PsGetCurrentProcessId() && NT_SUCCESS(PsLookupProcessByProcessId(CapturedCid.UniqueProcess, &Process)))
				{
					data = (svc_nt_open_process_data *)
						ExAllocatePoolWithTag(PagedPool, sizeof(svc_nt_open_process_data), 'TXSB');

					if (data)
					{
						RtlZeroMemory(data, sizeof(svc_nt_open_process_data));
						data->protocol = svc_nt_open_process;
						data->size = sizeof(svc_nt_open_process_data);
						data->time = PerfGetSystemTime();
						data->ProcessId = (ULONG)PsGetCurrentProcessId();
						data->TargetProcessId = (ULONG)PsGetProcessId(Process);
						data->DesiredAccess = (ULONG)DesiredAccess;
						data->ResultStatus = (ULONG)status;
						m_EventList->SendEvent(data);
					}
					ObDereferenceObject(Process);
				}
			}
		}
		__except (EXCEPTION_EXECUTE_HANDLER) 
		{
			if (data != NULL)
			{
				ExFreePoolWithTag(data, 'TXSB');
				data = NULL;
			}
		}
	}

	return status;
}

//NtTerminateProcess

NTSTATUS NTAPI NewNtTerminateProcess(
	_In_opt_ HANDLE ProcessHandle,
	_In_ NTSTATUS ExitStatus
);

static ShadowHookTarget m_NtTerminateProcessHookTarget =
{
	NewNtTerminateProcess,
	nullptr,
};

NTSTATUS NTAPI NewNtTerminateProcess(
	_In_opt_ HANDLE ProcessHandle,
	_In_ NTSTATUS ExitStatus
)
{
	const auto original = (NTSTATUS(NTAPI *)(HANDLE, NTSTATUS))m_NtTerminateProcessHookTarget.original_call;
	
	svc_nt_terminate_process_data *data = NULL;

	if (m_EventList->IsCapturing())
	{
		if (ARGUMENT_PRESENT(ProcessHandle) && ProcessHandle != NtCurrentProcess() && !ObIsKernelHandle(ProcessHandle) && ExGetPreviousMode() == UserMode)
		{
			PEPROCESS Process = NULL;
			HANDLE PrcoessId = NULL;

			if (NT_SUCCESS(GetProcessIdByHandle(ProcessHandle, &PrcoessId)) && NT_SUCCESS(PsLookupProcessByProcessId(PrcoessId, &Process)))
			{
				data = (svc_nt_terminate_process_data *)
					ExAllocatePoolWithTag(PagedPool, sizeof(svc_nt_terminate_process_data), 'TXSB');

				if (data)
				{
					RtlZeroMemory(data, sizeof(svc_nt_terminate_process_data));
					data->protocol = svc_nt_terminate_process;
					data->size = sizeof(svc_nt_terminate_process_data);
					data->time = PerfGetSystemTime();
					data->ProcessId = (ULONG)PsGetCurrentProcessId();
					data->TargetProcessId = (ULONG)PsGetProcessId(Process);
					data->ResultStatus = (ULONG)STATUS_SUCCESS;
				}
				ObDereferenceObject(Process);
			}
		}
	}

	NTSTATUS status = original(ProcessHandle, ExitStatus);

	if (data)
	{
		data->ResultStatus = (ULONG)status;
		m_EventList->SendEvent(data);
	}

	return status;
}

//NtReadVirtualMemory

NTSTATUS NTAPI NewNtReadVirtualMemory(
	__in HANDLE ProcessHandle,
	__in_opt PVOID BaseAddress,
	__out_bcount(BufferSize) PVOID Buffer,
	__in SIZE_T BufferSize,
	__out_opt PSIZE_T NumberOfBytesRead
);

static ShadowHookTarget m_NtReadVirtualMemoryHookTarget =
{
	NewNtReadVirtualMemory,
	nullptr,
};

NTSTATUS NTAPI NewNtReadVirtualMemory(
	__in HANDLE ProcessHandle,
	__in_opt PVOID BaseAddress,
	__out_bcount(BufferSize) PVOID Buffer,
	__in SIZE_T BufferSize,
	__out_opt PSIZE_T NumberOfBytesRead
)
{
	const auto original = (NTSTATUS(NTAPI *)(HANDLE, PVOID, PVOID, SIZE_T, PSIZE_T))m_NtReadVirtualMemoryHookTarget.original_call;

	NTSTATUS status = original(ProcessHandle, BaseAddress, Buffer, BufferSize, NumberOfBytesRead);
	
	svc_nt_read_virtual_mem_data *data = NULL;

	if (m_EventList->IsCapturing())
	{
		if (ARGUMENT_PRESENT(ProcessHandle) && ProcessHandle != NtCurrentProcess() && !ObIsKernelHandle(ProcessHandle) && ExGetPreviousMode() == UserMode)
		{
			PEPROCESS Process = NULL;
			HANDLE PrcoessId = NULL;
			if (NT_SUCCESS(GetProcessIdByHandle(ProcessHandle, &PrcoessId)) && PrcoessId != PsGetCurrentProcessId())
			{
				if (NT_SUCCESS(PsLookupProcessByProcessId(PrcoessId, &Process)))
				{
					data = (svc_nt_read_virtual_mem_data *)
						ExAllocatePoolWithTag(PagedPool, sizeof(svc_nt_read_virtual_mem_data), 'TXSB');

					if (data)
					{
						RtlZeroMemory(data, sizeof(svc_nt_read_virtual_mem_data));
						data->protocol = svc_nt_read_virtual_mem;
						data->size = sizeof(svc_nt_read_virtual_mem_data);
						data->time = PerfGetSystemTime();
						data->ProcessId = (ULONG)PsGetCurrentProcessId();
						data->TargetProcessId = (ULONG)PsGetProcessId(Process);
						data->BaseAddress = (ULONG64)BaseAddress;
						data->BufferSize = (ULONG64)BufferSize;
						data->ResultStatus = (ULONG)status;
						m_EventList->SendEvent(data);
					}
					ObDereferenceObject(Process);
				}				
			}
		}
	}

	return status;
}

//NtWriteVirtualMemory

NTSTATUS NTAPI NewNtWriteVirtualMemory(
	__in HANDLE ProcessHandle,
	__in_opt PVOID BaseAddress,
	__out_bcount(BufferSize) PVOID Buffer,
	__in SIZE_T BufferSize,
	__out_opt PSIZE_T NumberOfBytesRead
);

static ShadowHookTarget m_NtWriteVirtualMemoryHookTarget =
{
	NewNtWriteVirtualMemory,
	nullptr,
};

NTSTATUS NTAPI NewNtWriteVirtualMemory(
	__in HANDLE ProcessHandle,
	__in_opt PVOID BaseAddress,
	__out_bcount(BufferSize) PVOID Buffer,
	__in SIZE_T BufferSize,
	__out_opt PSIZE_T NumberOfBytesWritten
)
{
	const auto original = (NTSTATUS(NTAPI *)(HANDLE, PVOID, PVOID, SIZE_T, PSIZE_T))m_NtWriteVirtualMemoryHookTarget.original_call;

	NTSTATUS status = original(ProcessHandle, BaseAddress, Buffer, BufferSize, NumberOfBytesWritten);

	svc_nt_write_virtual_mem_data *data = NULL;

	if (m_EventList->IsCapturing())
	{
		if (ARGUMENT_PRESENT(ProcessHandle) && ProcessHandle != NtCurrentProcess() && !ObIsKernelHandle(ProcessHandle) && ExGetPreviousMode() == UserMode)
		{
			PEPROCESS Process = NULL;
			HANDLE PrcoessId = NULL;
			if (NT_SUCCESS(GetProcessIdByHandle(ProcessHandle, &PrcoessId)) && PrcoessId != PsGetCurrentProcessId())
			{
				if (NT_SUCCESS(PsLookupProcessByProcessId(PrcoessId, &Process)))
				{
					data = (svc_nt_write_virtual_mem_data *)
						ExAllocatePoolWithTag(PagedPool, sizeof(svc_nt_write_virtual_mem_data), 'TXSB');

					if (data)
					{
						RtlZeroMemory(data, sizeof(svc_nt_write_virtual_mem_data));
						data->protocol = svc_nt_write_virtual_mem;
						data->size = sizeof(svc_nt_write_virtual_mem_data);
						data->time = PerfGetSystemTime();
						data->ProcessId = (ULONG)PsGetCurrentProcessId();
						data->TargetProcessId = (ULONG)PsGetProcessId(Process);
						data->BaseAddress = (ULONG64)BaseAddress;
						data->BufferSize = (ULONG64)BufferSize;
						data->ResultStatus = (ULONG)status;
						m_EventList->SendEvent(data);
					}
					ObDereferenceObject(Process);
				}
			}
		}
	}

	return status;
}

//NtLoadDriver

NTSTATUS NTAPI NewNtLoadDriver(PUNICODE_STRING RegisterPath);

static ShadowHookTarget m_NtLoadDriverHookTarget =
{
	NewNtLoadDriver,
	nullptr,
};

NTSTATUS NTAPI NewNtLoadDriver(PUNICODE_STRING DriverServiceName)
{
	const auto original = (NTSTATUS(NTAPI*)(PUNICODE_STRING))m_NtLoadDriverHookTarget.original_call;

	svc_nt_load_driver_data *data = NULL;

	if (m_EventList->IsCapturing())
	{
		__try
		{
			UNICODE_STRING NewServiceName;
			ProbeForRead(DriverServiceName, sizeof(UNICODE_STRING), sizeof(ULONG));
			ProbeForRead(DriverServiceName->Buffer, DriverServiceName->Length, sizeof(WCHAR));

			data = (svc_nt_load_driver_data *)
				ExAllocatePoolWithTag(PagedPool, sizeof(svc_nt_load_driver_data), 'TXSB');
			if (data)
			{
				RtlZeroMemory(data, sizeof(svc_nt_load_driver_data));
				data->protocol = svc_nt_load_driver;
				data->size = sizeof(svc_nt_load_driver_data);
				data->time = PerfGetSystemTime();
				data->ProcessId = (ULONG)PsGetCurrentProcessId();
				data->ResultStatus = STATUS_SUCCESS;
				NewServiceName.Buffer = (PWCH)data->RegisterPath;;
				NewServiceName.Length = 0;
				NewServiceName.MaximumLength = sizeof(data->RegisterPath) - sizeof(WCHAR);
				RtlCopyUnicodeString(&NewServiceName, DriverServiceName);

				//Read ImagePath...
				HANDLE keyHandle = NULL;
				OBJECT_ATTRIBUTES oa;
				InitializeObjectAttributes(&oa, DriverServiceName,
					OBJ_CASE_INSENSITIVE | OBJ_KERNEL_HANDLE,
					0, 0);
				if (NT_SUCCESS(ZwOpenKey(&keyHandle, KEY_READ, &oa)))
				{
					KEY_VALUE_PARTIAL_INFORMATION info = { 0 };
					ULONG ulValueSize = 0;
					UNICODE_STRING valueName;
					RtlInitUnicodeString(&valueName, L"ImagePath");
					if (ZwQueryValueKey(keyHandle, &valueName, KeyValuePartialInformation, &info, sizeof(info), &ulValueSize) && ulValueSize > 0 && (info.Type == REG_SZ || info.Type == REG_EXPAND_SZ))
					{
						PKEY_VALUE_PARTIAL_INFORMATION infoBuffer = (PKEY_VALUE_PARTIAL_INFORMATION)ExAllocatePoolWithTag(NonPagedPool, ulValueSize, 'TXSB');
						if (infoBuffer != NULL)
						{
							if (NT_SUCCESS(ZwQueryValueKey(keyHandle, &valueName, KeyValuePartialInformation,
								infoBuffer, ulValueSize, &ulValueSize)))
							{
								UNICODE_STRING ustrSrcValue, ustrDstValue;
								ustrSrcValue.Buffer = (PWCH)infoBuffer->Data;
								ustrSrcValue.Length = ulValueSize;
								ustrSrcValue.MaximumLength = ulValueSize;
								ustrDstValue.Buffer = data->ImagePath;
								ustrDstValue.Length = 0;
								ustrDstValue.MaximumLength = sizeof(data->ImagePath) - sizeof(WCHAR);
								RtlCopyUnicodeString(&ustrDstValue, &ustrSrcValue);
							}
							ExFreePoolWithTag(infoBuffer, 'TXSB');
						}
					}

					ZwClose(keyHandle);
				}
			}
		}
		__except (EXCEPTION_EXECUTE_HANDLER)
		{
			if (data != NULL)
			{
				ExFreePoolWithTag(data, 'TXSB');
				data = NULL;
			}
		}
	}

	const auto status = original(DriverServiceName);

	if (data)
	{
		data->ResultStatus = status;
		m_EventList->SendEvent(data);
	}

	return status;
}

//free

VOID FreeShadowHooks(VOID)
{
#define FREE_SHADOWHOOK(a)if(m_##a##HookTarget.original_call != nullptr)\
	ExFreePoolWithTag(m_##a##HookTarget.original_call, kHyperPlatformCommonPoolTag);

	FREE_SHADOWHOOK(NtQuerySystemInformation)
	FREE_SHADOWHOOK(NtOpenProcess)
	FREE_SHADOWHOOK(NtTerminateProcess)
	FREE_SHADOWHOOK(NtReadVirtualMemory)
	FREE_SHADOWHOOK(NtWriteVirtualMemory)
	FREE_SHADOWHOOK(NtLoadDriver)
#undef FREE_SHADOWHOOK
}

NTSTATUS ShadowHookInitialization(SharedShadowHookData* shared_sh_data) 
{
	NTSTATUS status;
#define INSTALL_SHADOWHOOK(a) if (dynData.pfn##a)\
	{\
		if (!ShInstallHook(shared_sh_data, dynData.pfn##a, &m_##a##HookTarget))\
		{\
			FreeShadowHooks();\
			return STATUS_UNSUCCESSFUL;\
		}\
	}

	INSTALL_SHADOWHOOK(NtQuerySystemInformation)
	INSTALL_SHADOWHOOK(NtOpenProcess)
	INSTALL_SHADOWHOOK(NtTerminateProcess)
	INSTALL_SHADOWHOOK(NtReadVirtualMemory)
	INSTALL_SHADOWHOOK(NtWriteVirtualMemory)
	INSTALL_SHADOWHOOK(NtLoadDriver)
#undef INSTALL_SHADOWHOOK

	status = ShEnableHooks();
	if (!NT_SUCCESS(status))
	{
		FreeShadowHooks();
		return status;
	}

	return STATUS_SUCCESS;
}

VOID ShadowHookTermination(VOID)
{
	PAGED_CODE();

	ShDisableHooks();
	//UtilSleep(500);
	FreeShadowHooks();
}